"""quick_nrrd_2_niigz.py
    此脚本使用类似并行化写法，提高nrrd到niigz的转换效率
    使用时，修改main中的输入和输出位置即可

    不推荐嵌入到原先的代码中，可当做预先处理脚本，提高整体的代码执行速度
"""

import os
import time
import traceback
from omp4py import omp
import SimpleITK as sitk

def _nrrd_2_nii_impl(input_nrrd, output_niigz):
    if os.path.exists(output_niigz):
        return True, ""

    try:
        sitk_image = sitk.ReadImage(input_nrrd)
        sitk.WriteImage(sitk_image, output_niigz)
        return True, ""
    except Exception as e:
        return False, traceback.format_exc()

@omp
def _nrrd_2_nii_omp(input_folder, output_folder):
    omp_dict = {}
    for root, dir, files in os.walk(input_folder):
        with omp("parallel"):
            for each_file in files:
                input_path = os.path.join(root, each_file)
                output_path = os.path.join(output_folder, each_file.replace('.nrrd', '.nii.gz'))
                omp_dict[each_file] = _nrrd_2_nii_impl(input_path, output_path)
                print("Now omp process {} to niigz".format(each_file))

    ok_flag = True
    for each_key in omp_dict:
        status, msg = omp_dict[each_key]
        if not status:
            ok_flag = False
            print("File {} convert to niigz failed, exception is: {}".format(each_key, msg))

    return ok_flag, omp_dict

def _nrrd_2_nii_cpu(input_folder, output_folder):
    cpu_dict = {}
    ok_flag = True
    for root, dir, files in os.walk(input_folder):
        for each_file in files:
            input_path = os.path.join(root, each_file)
            output_path = os.path.join(output_folder, each_file.replace('.nrrd', '.nii.gz'))
            status, msg = _nrrd_2_nii_impl(input_path, output_path)
            cpu_dict[each_file] = (status, msg)
            print("Now process {} to niigz".format(each_file))

            if not status:
                ok_flag = False
                print("File {} convert to niigz failed, exception is: {}".format(each_file, msg))

    return ok_flag, cpu_dict


@omp
def nrrd_2_nii(input_folder, output_folder, use_openmp: bool = False):
    os.makedirs(output_folder, exist_ok=True)

    if use_openmp:
        return _nrrd_2_nii_omp(input_folder, output_folder)
    else:
        return _nrrd_2_nii_cpu(input_folder, output_folder)


if __name__ == '__main__':
    start = time.perf_counter()
    print(nrrd_2_nii("../SampleData/7", "../SampleData/niigz", use_openmp=True))
    end = time.perf_counter()
    print("time cost", end - start)